using RAG.Domain.Services;
using System.Text;

namespace RAG.Infrastructure.Services;

public class TextChunkingService : ITextChunkingService
{
    public Task<List<string>> ChunkTextAsync(string text, int chunkSize = 1000, int overlap = 200)
    {
        var chunks = new List<string>();

        if (string.IsNullOrWhiteSpace(text))
            return Task.FromResult(chunks);

        // 按段落分割
        var paragraphs = text.Split(new[] { "\n\n", "\r\n\r\n" }, StringSplitOptions.RemoveEmptyEntries);

        var currentChunk = new StringBuilder();

        foreach (var paragraph in paragraphs)
        {
            // 如果当前段落加上现有chunk超过了大小限制
            if (currentChunk.Length + paragraph.Length > chunkSize && currentChunk.Length > 0)
            {
                chunks.Add(currentChunk.ToString().Trim());

                // 保留重叠内容
                var overlapText = GetOverlapText(currentChunk.ToString(), overlap);
                currentChunk.Clear();
                currentChunk.Append(overlapText);
            }

            currentChunk.AppendLine(paragraph);

            // 如果单个段落就超过了chunk大小，需要进一步分割
            if (paragraph.Length > chunkSize)
            {
                var sentences = SplitIntoSentences(paragraph);
                currentChunk.Clear();

                foreach (var sentence in sentences)
                {
                    if (currentChunk.Length + sentence.Length > chunkSize && currentChunk.Length > 0)
                    {
                        chunks.Add(currentChunk.ToString().Trim());

                        var overlapText = GetOverlapText(currentChunk.ToString(), overlap);
                        currentChunk.Clear();
                        currentChunk.Append(overlapText);
                    }

                    currentChunk.Append(sentence + " ");
                }
            }
        }

        // 添加最后一个chunk
        if (currentChunk.Length > 0)
        {
            chunks.Add(currentChunk.ToString().Trim());
        }

        return Task.FromResult(chunks);
    }

    private string GetOverlapText(string text, int overlapSize)
    {
        if (text.Length <= overlapSize)
            return text;

        var lastPart = text.Substring(text.Length - overlapSize);

        // 尝试在词边界处截断
        var lastSpaceIndex = lastPart.LastIndexOf(' ');
        if (lastSpaceIndex > overlapSize / 2)
        {
            return lastPart.Substring(lastSpaceIndex + 1);
        }

        return lastPart;
    }

    private List<string> SplitIntoSentences(string text)
    {
        var sentences = new List<string>();
        var sentenceEnders = new[] { '.', '!', '?', '。', '！', '？' };

        var currentSentence = new StringBuilder();

        for (int i = 0; i < text.Length; i++)
        {
            currentSentence.Append(text[i]);

            if (sentenceEnders.Contains(text[i]))
            {
                // 检查是否是句子结尾（后面是空格或文本结束）
                if (i == text.Length - 1 || char.IsWhiteSpace(text[i + 1]))
                {
                    sentences.Add(currentSentence.ToString().Trim());
                    currentSentence.Clear();
                }
            }
        }

        // 添加剩余内容
        if (currentSentence.Length > 0)
        {
            sentences.Add(currentSentence.ToString().Trim());
        }

        return sentences.Where(s => !string.IsNullOrWhiteSpace(s)).ToList();
    }
}
